import os
import json
import faiss
import numpy as np
from typing import List, Optional
from pydantic import BaseModel
from fastapi import FastAPI, HTTPException
from sentence_transformers import SentenceTransformer
import uvicorn

# === 配置 ===
PARAGRAPH_FILE = "wiki_paragraphs.tsv"
TITLE_MAP_FILE = "wiki_titles.json"
FAISS_INDEX_FILE = "wiki_paragraphs.index"
EMBEDDING_MODEL = "/share/models/BAAI/bge-base-en-v1.5"


# === 数据模型 ===
class SearchRequest(BaseModel):
    query: str
    top_k: int = 10
    threshold: Optional[float] = None


class SearchResult(BaseModel):
    doc_id: int
    title: str
    content: str
    score: float


class SearchResponse(BaseModel):
    query: str
    results: List[SearchResult]
    total_found: int


# 兼容原有调用方式的响应格式
class CompatibleResponse(BaseModel):
    response_code: int
    data: dict
    message: Optional[str] = None


# === 检索服务类 ===
class WikipediaSearchService:
    def __init__(self):
        self.model = None
        self.index = None
        self.paragraphs = {}
        self.titles = {}
        self.is_loaded = False

    def load_resources(self):
        """加载模型和索引"""
        try:
            print("[*] 正在加载嵌入模型...")
            self.model = SentenceTransformer(EMBEDDING_MODEL)

            print("[*] 正在加载FAISS索引...")
            if not os.path.exists(FAISS_INDEX_FILE):
                raise FileNotFoundError(f"FAISS索引文件不存在: {FAISS_INDEX_FILE}")
            self.index = faiss.read_index(FAISS_INDEX_FILE)

            print("[*] 正在加载段落数据...")
            if not os.path.exists(PARAGRAPH_FILE):
                raise FileNotFoundError(f"段落文件不存在: {PARAGRAPH_FILE}")

            with open(PARAGRAPH_FILE, "r", encoding="utf-8") as f:
                for line in f:
                    parts = line.strip().split("\t", 1)
                    if len(parts) == 2:
                        doc_id, content = parts
                        self.paragraphs[int(doc_id)] = content

            print("[*] 正在加载标题映射...")
            if not os.path.exists(TITLE_MAP_FILE):
                raise FileNotFoundError(f"标题映射文件不存在: {TITLE_MAP_FILE}")

            with open(TITLE_MAP_FILE, "r", encoding="utf-8") as f:
                titles_data = json.load(f)
                # 转换键为整数
                self.titles = {int(k): v for k, v in titles_data.items()}

            self.is_loaded = True
            print(f"[✔] 资源加载完成！共加载 {len(self.paragraphs)} 个段落")

        except Exception as e:
            print(f"[✗] 加载资源失败: {str(e)}")
            raise e

    def search(self, query: str, top_k: int = 10, threshold: Optional[float] = None) -> List[SearchResult]:
        """执行检索"""
        if not self.is_loaded:
            raise RuntimeError("服务未初始化，请先加载资源")

        # 对查询进行编码
        query_embedding = self.model.encode([query], normalize_embeddings=True)

        # 在FAISS索引中搜索
        scores, doc_ids = self.index.search(query_embedding, top_k)

        results = []
        for i, (score, doc_id) in enumerate(zip(scores[0], doc_ids[0])):
            # 跳过无效的doc_id
            if doc_id == -1:
                continue

            # 应用阈值过滤
            if threshold is not None and score < threshold:
                continue

            # 获取内容和标题
            content = self.paragraphs.get(doc_id, "内容未找到")
            title = self.titles.get(doc_id, "标题未知")

            results.append(SearchResult(
                doc_id=int(doc_id),
                title=title,
                content=content,
                score=float(score)
            ))

        return results


# === 创建FastAPI应用 ===
app = FastAPI(
    title="Wikipedia检索API",
    description="基于FAISS和BGE嵌入的Wikipedia段落检索服务",
    version="1.0.0"
)

# 全局服务实例
search_service = WikipediaSearchService()


@app.on_event("startup")
async def startup_event():
    """应用启动时加载资源"""
    try:
        search_service.load_resources()
    except Exception as e:
        print(f"启动失败: {str(e)}")
        raise e


@app.get("/status")
async def root():
    """健康检查端点（保持原有根路径功能）"""
    return {
        "message": "Wikipedia检索API运行中",
        "status": "healthy",
        "loaded": search_service.is_loaded
    }


@app.get("/health")
async def health_check():
    """详细健康检查"""
    return {
        "status": "healthy" if search_service.is_loaded else "not_ready",
        "paragraphs_count": len(search_service.paragraphs),
        "titles_count": len(search_service.titles),
        "model_loaded": search_service.model is not None,
        "index_loaded": search_service.index is not None
    }


@app.post("/search", response_model=SearchResponse)
async def search_paragraphs(request: SearchRequest):
    """检索段落（标准格式）"""
    if not search_service.is_loaded:
        raise HTTPException(status_code=503, detail="服务未就绪，请等待资源加载完成")

    if not request.query.strip():
        raise HTTPException(status_code=400, detail="查询不能为空")

    if request.top_k < 1 or request.top_k > 100:
        raise HTTPException(status_code=400, detail="top_k必须在1-100之间")

    try:
        results = search_service.search(
            query=request.query,
            top_k=request.top_k,
            threshold=request.threshold
        )

        return SearchResponse(
            query=request.query,
            results=results,
            total_found=len(results)
        )

    except Exception as e:
        raise HTTPException(status_code=500, detail=f"检索失败: {str(e)}")


@app.post("/", response_model=CompatibleResponse)
async def compatible_search(request: SearchRequest):
    """兼容原有调用方式的检索端点"""
    try:
        if not search_service.is_loaded:
            return CompatibleResponse(
                response_code=503,
                data={},
                message="服务未就绪，请等待资源加载完成"
            )

        if not request.query.strip():
            return CompatibleResponse(
                response_code=400,
                data={},
                message="查询不能为空"
            )

        if request.top_k < 1 or request.top_k > 100:
            return CompatibleResponse(
                response_code=400,
                data={},
                message="top_k必须在1-100之间"
            )

        results = search_service.search(
            query=request.query,
            top_k=request.top_k,
            threshold=request.threshold
        )

        # 转换为兼容格式
        contexts = []
        for result in results:
            contexts.append({
                "doc_id": result.doc_id,
                "title": result.title,
                "content": result.content,
                "score": result.score
            })

        return CompatibleResponse(
            response_code=200,
            data={"context": contexts}
        )

    except Exception as e:
        return CompatibleResponse(
            response_code=500,
            data={},
            message=f"检索失败: {str(e)}"
        )


@app.get("/search")
async def search_get(
        q: str,
        top_k: int = 10,
        threshold: Optional[float] = None
):
    """GET方式检索（便于测试）"""
    request = SearchRequest(query=q, top_k=top_k, threshold=threshold)
    return await search_paragraphs(request)


@app.get("/stats")
async def get_stats():
    """获取统计信息"""
    if not search_service.is_loaded:
        raise HTTPException(status_code=503, detail="服务未就绪")

    return {
        "total_paragraphs": len(search_service.paragraphs),
        "unique_titles": len(set(search_service.titles.values())),
        "index_dimension": search_service.index.d if search_service.index else 0,
        "model_name": EMBEDDING_MODEL
    }


# === 运行服务 ===
if __name__ == "__main__":
    # 确保必要文件存在
    required_files = [PARAGRAPH_FILE, TITLE_MAP_FILE, FAISS_INDEX_FILE]
    missing_files = [f for f in required_files if not os.path.exists(f)]

    if missing_files:
        print(f"[✗] 缺少必要文件: {missing_files}")
        print("[*] 请先运行索引构建脚本生成这些文件")
        exit(1)

    print("[*] 启动Wikipedia检索API服务...")
    uvicorn.run(
        app,
        host="0.0.0.0",
        port=8088,  # 默认使用8088端口
        log_level="info"
    )